package main

import (
	"log"
	"net"
	"sync"
	"time"
)

// ModbusProxy 结构体
type ModbusProxy struct {
	listenAddr string   // 监听地址
	targetAddr string   // 目标地址
	timeout    int      // 超时时间(秒)
	mu         sync.Mutex
}

// NewModbusProxy 创建新的Modbus代理
func NewModbusProxy(listenAddr, targetAddr string, timeout int) *ModbusProxy {
	return &ModbusProxy{
		listenAddr: listenAddr,
		targetAddr: targetAddr,
		timeout:    timeout,
	}
}

// Start 启动代理服务
func (p *ModbusProxy) Start() error {
	listener, err := net.Listen("tcp", p.listenAddr)
	if err != nil {
		return err
	}
	defer listener.Close()

	log.Printf("Modbus proxy started, listening on %s, forwarding to %s\n", p.listenAddr, p.targetAddr)

	for {
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("Accept error: %v\n", err)
			continue
		}

		go p.handleConnection(clientConn)
	}
}

// handleConnection 处理客户端连接
func (p *ModbusProxy) handleConnection(clientConn net.Conn) {
	defer clientConn.Close()

	// 连接目标设备
	targetConn, err := net.DialTimeout("tcp", p.targetAddr, time.Duration(p.timeout)*time.Second)
	if err != nil {
		log.Printf("Failed to connect to target %s: %v\n", p.targetAddr, err)
		return
	}
	defer targetConn.Close()

	log.Printf("Connection established: %s <-> %s\n", clientConn.RemoteAddr(), p.targetAddr)

	// 双向转发数据
	var wg sync.WaitGroup
	wg.Add(2)

	// 客户端 -> 目标设备
	go func() {
		defer wg.Done()
		p.forwardData(clientConn, targetConn)
	}()

	// 目标设备 -> 客户端
	go func() {
		defer wg.Done()
		p.forwardData(targetConn, clientConn)
	}()

	wg.Wait()
}

// forwardData 转发数据
func (p *ModbusProxy) forwardData(src, dst net.Conn) {
	buf := make([]byte, 4096)
	for {
		n, err := src.Read(buf)
		if err != nil {
			log.Printf("Read error from %s: %v\n", src.RemoteAddr(), err)
			break
		}

		p.mu.Lock()
		_, err = dst.Write(buf[:n])
		p.mu.Unlock()
		
		if err != nil {
			log.Printf("Write error to %s: %v\n", dst.RemoteAddr(), err)
			break
		}
	}
}

func main() {
	// 配置参数
	listenAddr := ":5020"       // 代理监听地址和端口
	targetAddr := "192.168.1.100:502" // 目标Modbus设备地址
	timeout := 5               // 连接超时时间(秒)

	proxy := NewModbusProxy(listenAddr, targetAddr, timeout)
	log.Fatal(proxy.Start())
}